/**
 * Copyright 2019-2020 Huawei Technologies Co., Ltd
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#ifndef FUSION_ENGINE_OPTIMIZER_GRAPH_OPTIMIZER_UB_FUSION_REACHABILITY_MAP_H_
#define FUSION_ENGINE_OPTIMIZER_GRAPH_OPTIMIZER_UB_FUSION_REACHABILITY_MAP_H_
#include <list>
#include <memory>
#include <string>
#include <unordered_map>
#include <vector>
#include "common/fe_log.h"
#include "common/fe_utils.h"
#include "common/math_util.h"
#include "common/scope_allocator.h"
#include "common/util/constants.h"
#include "common/util/op_info_util.h"
#include "graph/debug/ge_attr_define.h"
#include "graph/node.h"

namespace fe {
class ReachabilityMap {
 public:
  /* Sets up a graph with no edges and where the nodes correspond to the given
   * instructions. */
  explicit ReachabilityMap(ge::ComputeGraph &graph);

  ~ReachabilityMap();

  /* Computes and returns the reachability between two nodes in the
   * computation. The returned ReachabilityMap is constructed such that
   * ReachabilityMap::IsReachable(a, b) returns true iff there exists a
   * directed path (from producer to consumer) from 'a' to 'b'. Only data
   * operands is considered for reachability. Trivially an instruction is
   * reachable from itself. */
  static std::shared_ptr<ReachabilityMap> Build(ge::ComputeGraph &graph);

  // Returns true if "b" is reachable from "a"
  // Note that this function only correctly answers queries about reachability
  // if the set of edges that have been provided to this class are transitive.
  bool IsReachable(const ge::NodePtr &a, const ge::NodePtr &b) const;

  // update reachablity map for fused nodes.
  void Update(ge::ComputeGraph &graph, vector<ge::NodePtr> &fusion_nodes);

 private:
  // A bit-vector implementation specialized for this use case which provides a
  // fast bitwise OR operation not available in tensorflow::gtl::BitMap.
  class BitVector {
   public:
    BitVector() : size_(0) {}

    ~BitVector() = default;

    explicit BitVector(size_t size) : size_(size), vector_((size + k_bits - 1) / k_bits, 0) {}

    // Return the bit at the given index.
    bool Get(size_t index) const {
      if (index >= 0 && index < size_) {
        return vector_[index / k_bits] & (1ull << (index % k_bits));
      } else {
        FE_LOGE("index %u is not valid. Total size is %u", index, size_);
        return false;
      }
    }

    // Set the bit at the given index.
    void Set(size_t index) {
      if (index >= 0 && index < size_) {
        vector_[index / k_bits] |= 1ull << (index % k_bits);
      } else {
        FE_LOGE("index %u is not valid. Total size is %u", index, size_);
        return;
      }
    }

    // Set this bitvector to the Logical OR of this bitvector and 'other'.
    void OrWith(const BitVector &other) {
      for (size_t i = 0; i < vector_.size(); ++i) {
        vector_[i] |= other.vector_[i];
      }
    }

    // Set the bitvector to all zeros.
    void SetToZero() { std::fill(vector_.begin(), vector_.end(), 0); }

    bool operator==(const BitVector &other) const { return vector_ == other.vector_; }

    bool operator!=(const BitVector &other) const { return vector_ != other.vector_; }

   private:
    using Word = uint64_t;
    static constexpr size_t k_bits = 64;

    // Number of bits in the bitvector.
    size_t size_;

    std::vector<Word> vector_;
  };

  int64_t GetIndex(const ge::NodePtr &node) const {
    auto iter = name_to_index_.find(node->GetName());
    if (iter != name_to_index_.end()) {
      return iter->second;
    } else {
      FE_LOGW("node %s is not found in name_to_index_", node->GetName().c_str());
      return 0;
    }
  }

  // Return the bitvector storing the reachability-to of the given instruction.
  const BitVector &GetBitVector(const ge::NodePtr &node) const { return bit_vectors_[GetIndex(node)]; }

  BitVector &GetBitVector(const ge::NodePtr &node) { return bit_vectors_[GetIndex(node)]; }

  // The number of instructions in the reachability map.
  const size_t size_;

  // Dense assignment from node::GetOpDesc()::Id to number. These numbers
  // index into the bit_vectors_ vector and into the bits within a BitVector.
  std::unordered_map<uint64_t, int> indices_;

  // Bitvectors holding the reachability to each instruction. The bit vector for
  // instruction X includes ones for each instruction which X is reachable from.
  std::vector<BitVector> bit_vectors_;

  /* map of op name and op index in this class.
   * op index starts from 0. */
  std::map<std::string, int64_t> name_to_index_;

  void SetReachabilityToUnionHelper(const std::vector<ge::NodePtr> &inputs, const ge::NodePtr &node,
                                    BitVector *bit_vector);

  void FastSetReachabilityToUnion(const std::vector<ge::NodePtr> &inputs, const ge::NodePtr &node);
};
}
#endif  // FUSION_ENGINE_OPTIMIZER_GRAPH_OPTIMIZER_UB_FUSION_REACHABILITY_MAP_H_
